package CanvaWeb.CanvaWeb.WebSocket;

import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.ArrayList;

@Component
public class WebSocketHandler extends TextWebSocketHandler {

    private final Map<Long, List<WebSocketSession>> roomSessions = new HashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long roomId = extractRoomIdFromSession(session);

        roomSessions.computeIfAbsent(roomId, k -> new ArrayList<>()).add(session);

        System.out.println("WebSocket connection established to room: " + roomId);
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        Long roomId = extractRoomIdFromSession(session);

        // 将收到的消息广播到房间内的所有客户端
        List<WebSocketSession> sessions = roomSessions.get(roomId);
        if (sessions != null) {
            for (WebSocketSession s : sessions) {
                if (s.isOpen()) {
                    s.sendMessage(new TextMessage(message.getPayload()));
                }
            }
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Long roomId = extractRoomIdFromSession(session);

        // 从房间中移除关闭连接的 session
        List<WebSocketSession> sessions = roomSessions.get(roomId);
        if (sessions != null) {
            sessions.remove(session);
            System.out.println("WebSocket connection closed from room: " + roomId);
        }
    }

    private Long extractRoomIdFromSession(WebSocketSession session) {
        String uri = session.getUri().toString();
        String roomId = uri.substring(uri.lastIndexOf("/") + 1); // 获取路径中的最后部分作为 roomId
        return Long.parseLong(roomId);
    }

    // 添加获取房间 WebSocket 会话的公共方法
    public List<WebSocketSession> getSessionsByRoomId(Long roomId) {
        return roomSessions.getOrDefault(roomId, new ArrayList<>());
    }
}

